from datetime import datetime, timedelta
from typing import Union

import pandas as pd


def truncate_date(date: Union[str, datetime],
                  bucket_width: str,
                  offset: int = 0) -> datetime:
    """
    将日期截断到指定宽度的网格中

    参数:
        date: 要截断的日期，可以是字符串或datetime对象
        bucket_width: 桶的宽度，格式为数字加单位(如"1D", "1M", "1Y")
        offset: 锚点的偏移量，单位与bucket_width一致

    返回:
        截断后的datetime对象
    """
    # 解析输入日期
    if isinstance(date, str):
        date = pd.to_datetime(date)

    # 解析bucket_width
    try:
        width, unit = int(bucket_width[:-1]), bucket_width[-1].upper()
    except (ValueError, IndexError):
        raise ValueError(f"无效的bucket_width格式: {bucket_width}，应为如'1D', '1M'的格式")

    # 确定锚点基准日期
    if unit in ['M', 'Y']:  # 月或年单位
        anchor_base = datetime(2000, 1, 1)
    else:  # 天或更细单位
        anchor_base = datetime(2000, 1, 3)  # 2000-01-03是星期一

    # 应用偏移量到锚点
    anchor = apply_offset(anchor_base, width, unit, offset)

    # 根据单位执行不同的截断逻辑
    if unit == 'D':
        # 按天截断，计算距离锚点的天数
        days_diff = (date - anchor).days
        truncated = anchor + timedelta(days=days_diff // width * width)
    elif unit == 'M':
        # 按月截断，计算距离锚点的月数
        months_diff = (date.year - anchor.year) * 12 + (date.month - anchor.month)
        truncated_months = anchor.month + (months_diff // width) * width
        truncated_year = anchor.year + (truncated_months - 1) // 12
        truncated_month = (truncated_months - 1) % 12 + 1
        # 确保日期在新月份的有效范围内
        truncated = datetime(truncated_year, truncated_month,
                             min(date.day, get_days_in_month(truncated_year, truncated_month)))
    elif unit == 'Y':
        # 按年截断，计算距离锚点的年数
        years_diff = date.year - anchor.year
        truncated_year = anchor.year + (years_diff // width) * width
        truncated = datetime(truncated_year, date.month, min(date.day, get_days_in_month(truncated_year, date.month)))
    else:
        raise ValueError(f"不支持的时间单位: {unit}，仅支持'D', 'M', 'Y'")

    return truncated


def apply_offset(anchor: datetime, width: int, unit: str, offset: int) -> datetime:
    """将偏移量应用到锚点"""
    if unit == 'D':
        return anchor + timedelta(days=offset * width)
    elif unit == 'M':
        months = offset * width
        year, month = anchor.year + (anchor.month + months - 1) // 12, (anchor.month + months - 1) % 12 + 1
        day = min(anchor.day, get_days_in_month(year, month))
        return datetime(year, month, day)
    elif unit == 'Y':
        return datetime(anchor.year + offset * width, anchor.month,
                        min(anchor.day, get_days_in_month(anchor.year + offset * width, anchor.month)))
    else:
        raise ValueError(f"不支持的时间单位: {unit}，仅支持'D', 'M', 'Y'")


def get_days_in_month(year: int, month: int) -> int:
    """获取指定年月的天数"""
    if month in [4, 6, 9, 11]:
        return 30
    elif month == 2:
        return 29 if is_leap_year(year) else 28
    else:
        return 31


def is_leap_year(year: int) -> bool:
    """判断是否为闰年"""
    return year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)


if __name__ == '__main__':
    # 示例1: 天单位截断 (2000-01-03是星期一)
    date1 = "2020-01-05"
    bucket_width1 = "27D"  # 7天为一个桶
    offset1 = 0
    truncated1 = truncate_date(date1, bucket_width1, offset1)
    print(f"日期 {date1} 截断到{bucket_width1}网格(偏移{offset1})的结果: {truncated1}")

    # 示例2: 月单位截断
    date2 = "2020-05-15"
    bucket_width2 = "3M"  # 3个月为一个桶
    offset2 = 2
    truncated2 = truncate_date(date2, bucket_width2, offset2)
    print(f"日期 {date2} 截断到{bucket_width2}网格(偏移{offset2})的结果: {truncated2}")

    # 示例3: 年单位截断
    date3 = "2021-07-20"
    bucket_width3 = "5Y"  # 5年为一个桶
    offset3 = 1
    truncated3 = truncate_date(date3, bucket_width3, offset3)
    print(f"日期 {date3} 截断到{bucket_width3}网格(偏移{offset3})的结果: {truncated3}")
